from os import getcwd, path
from numpy import linspace, abs, zeros, real, arange, angle, sum ,sqrt, pi, unwrap, gradient, \
                  polyfit, polyval  
from numpy.fft import fft, fftfreq, fftshift
from enthought.traits.api import Any, Bool, Int, Float, Str, Range, Button, HasTraits, \
                                DelegatesTo, Instance, Enum, Event, Trait, Tuple, \
                                Callable, Array, List, Property, on_trait_change
from enthought.traits.ui.api import View, Item, Group, HGroup, VGroup, VGrid, \
                                    RangeEditor, HSplit, VFold, HFlow
from enthought.traits.ui.editors import TextEditor
from enthought.traits.ui.menu import OKButton, CancelButton
from enthought.enable.api import BaseTool, KeySpec, ColorTrait, LineStyle, \
                                 Component, ComponentEditor, Window
from enthought.chaco.api import Plot, ArrayPlotData, DataRange1D, \
                                ArrayDataSource, LinearMapper, BasePlotContainer, \
                                OverlayPlotContainer, VPlotContainer, LinePlot, DataLabel, \
                                create_line_plot, create_scatter_plot, PlotLabel, \
                                add_default_axes
from enthought.chaco.tools.api import PanTool, SimpleZoom, RangeSelection, RangeSelectionOverlay
from gprViewerUIs import wxOpenFile


class TraceViewer(HasTraits):
    # Load Panel Items
    file = Str()
    dir = Str()
    browse = Button(label='...')
    sample_min = Int(0, label="Sample min/max")
    sample_max = Int(-1)
    trace_no = Int(0)
    trace_max = Int(-1, label="Max:")
    show_button = Button(label="Show")
    show = Bool(False)
    # Processing Panel Items
    spectrum_button = Button(label="Spectrum")
    cspectrum_button = Button(label="Complex Spectrum")
    # Status Panel
    status_txt = Str("TraceViewer started ...")
    # Plot Panel Items
    container = Instance(BasePlotContainer)
    rangeselect = Instance(RangeSelection)
    
    load_panel = VGrid(
                       Item(name='file'), Item(name='browse', show_label=False),
                       Item(name='sample_min'), Item(name='sample_max', show_label=False),
                       Item(name='trace_no'), Item(name='trace_max', style='readonly'),
                       Item(name='show_button', show_label=False),
                       show_border=True, label="Choose data file and range")
    
    proc_panel = VGroup(
                        Item(name='spectrum_button', show_label=False, width=300),
                        Item(name='cspectrum_button', show_label=False),
                        show_border=True
                        )
    
    status_panel = VGroup(
                          Item(name='status_txt', show_label=False, editor=TextEditor(),
                               style='readonly',
                               full_size=False, height=200),
                          show_border=True, label="Status"
                          )
    
    view = View(HSplit(
                       VGroup(
                              load_panel, proc_panel, status_panel
                              ),
#                      '_',
                      Item(name="container", show_label=False, editor=ComponentEditor())
#                          visible_when="show==True")),
                       ),
                icon=r"../src/recources/icons/edit-32.png",
                title="Trace Viewer",
                width=800, height=0.9,
                resizable=True)
    
    def __init__(self, file="", dir=getcwd()):
        self.dir = dir
        self.file = file
        
    def _create_traceplot(self):
        x = self.x
        y = self.y       
        traceplot = create_line_plot((x, y), index_bounds = None, value_bounds = None,
                         orientation = 'v', color = 'blue', width = 1.0, dash = 'solid',
                         value_mapper_class = LinearMapper,
                         bgcolor = 'transparent', border_visible = True,
                         add_grid = False, add_axis = False, index_sort = 'ascending')
        add_default_axes(traceplot, orientation = 'flipped', vtitle = 'Time [ns]', htitle = 'Signal')
        traceplot.origin = "top left"
        traceplot.tools.append(PanTool(traceplot, drag_button="right"))
        zoom = SimpleZoom(component=traceplot, tool_mode="box", drag_button="left", always_on=True)
        traceplot.overlays.append(zoom)
        
        container = OverlayPlotContainer(padding=40, padding_left=60)
        container.add(traceplot)
        self.container = container
           
    def _create_spectrumplot(self, x, y):
        spectrumplot = create_line_plot((x, y), index_bounds = None, value_bounds = None,
                         orientation = 'v', color = 'green', width = 1.0, dash = 'solid',
                         value_mapper_class = LinearMapper,
                         bgcolor = 'transparent', border_visible = True,
                         add_grid = False, add_axis = False, index_sort = 'ascending')
        add_default_axes(spectrumplot, orientation = 'flipped', vtitle = 'Frequency [MHz]', htitle = 'Amplitude')
        spectrumplot.origin = "top left"
        spectrumplot.tools.append(PanTool(spectrumplot, drag_button="right"))
        zoom = SimpleZoom(component=spectrumplot, tool_mode="box", drag_button="left", always_on=True)
        spectrumplot.overlays.append(zoom)
        
        container = OverlayPlotContainer(padding=40, padding_left=60)
        container.add(spectrumplot)
        self.container = container
        
    def _create_complexspectrumplot(self,x, y1, y2):      
        amplitudeplot = create_line_plot((x, y1), index_bounds = None, value_bounds = None,
                         orientation = 'h', color = 'green', width = 1.0, dash = 'solid',
                         value_mapper_class = LinearMapper,
                         bgcolor = 'white', border_visible = True,
                         add_grid = False, add_axis = False, index_sort = 'ascending')
        add_default_axes(amplitudeplot, orientation = 'normal', htitle = 'Frequency [MHz]', vtitle = 'Amplitude')

        amplitudeplot.tools.append(PanTool(amplitudeplot, drag_button="right"))
        zoom = SimpleZoom(component=amplitudeplot, tool_mode="box", drag_button="left", always_on=True)
        amplitudeplot.overlays.append(zoom)
        
        phaseplot = create_line_plot((x, y2), index_bounds = None, value_bounds = None,
                         orientation = 'h', color = 'red', width = 1.0, dash = 'solid',
                         value_mapper_class = LinearMapper,
                         bgcolor = 'white', border_visible = True,
                         add_grid = False, add_axis = False, index_sort = 'ascending')
        add_default_axes(phaseplot, orientation = 'normal', htitle = 'Frequency [MHz]', vtitle = 'Phase')
        
        phaseplot.tools.append(PanTool(phaseplot, drag_button="right"))
#        zoom = SimpleZoom(component=phaseplot, tool_mode="box", drag_button="left", always_on=True)
#        phaseplot.overlays.append(zoom)
        self.rangeselect = RangeSelection(phaseplot, left_button_selects = True)
        self.rangeselect.on_trait_change(self.phase_center, "selection_completed")
        phaseplot.active_tool = self.rangeselect
        phaseplot.overlays.append(RangeSelectionOverlay(component=phaseplot))
        
        
        container = VPlotContainer(padding=40, padding_left=60, spacing=40)
        container.add(phaseplot)
        container.add(amplitudeplot)
        self.container = container
        
    def _file_changed(self):
        try:
            from tables import openFile, NoSuchNodeError
            h5file = openFile(path.join(self.dir,self.file), mode = 'r')
            try:
                sample_max, trace_max = h5file.root.data.traces.shape
                self.sample_max = int(sample_max)
                self.trace_max = int(trace_max)
            except NoSuchNodeError:
                self.sample_max, self.trace_max = h5file.root.profile.traces.shape
            h5file.close()
        except IOError:
            print "The chosen file has an unknown format."
    
    def _browse_fired(self):
        file, dir = wxOpenFile(self.dir, multi=False)
        self.dir = dir
        self.file = file
        
    def _show_button_fired(self):
        from processGPR import HUBRA
        gpr = HUBRA()
        gpr.readH5(path.join(self.dir,self.file),
                   self.trace_no, self.trace_no+1,
                   self.sample_min, self.sample_max)
        gpr.removeGating()
#        gpr.compensateSpread(a=2, power=False)
        self.x = gpr.t
        self.y = gpr.data[:,0]
        self.r = gpr.r
        self.deltaT = gpr.deltaT
        self._create_traceplot()
        self.status_txt = "%s\nTrace %i from file %s displayed" % (self.status_txt, self.trace_no, self.file) 
        
    def _spectrum_button_fired(self):
        Nfft = self.y.size
        freq = fftfreq(Nfft, self.deltaT)
        Y = abs(fft(self.y))
        x = freq[0:Nfft/2]*1e-6
        y = Y[0:Nfft/2]
        self._create_spectrumplot(x, y)
        
    def _cspectrum_button_fired(self):
        Nfft = self.y.size
        freq = fftfreq(Nfft, self.deltaT)
        c0 = 299792458.0
#        k = 2*pi*freq[0:Nfft/2]/(c0/sqrt(1.8))
        Y = fft(self.y)
        f = freq[0:Nfft/2]*1e-6 # [MHz]
        y1 = abs(Y)[0:Nfft/2]
        mag = y1/y1.max() # normalized amplitude spectrum
        pha = unwrap(angle(Y)[0:Nfft/2])
        
        self.freq = f
        self.pha = pha
        self._create_complexspectrumplot(f, mag, pha)
    
    def _get_index(self):
        smaller = self.freq>self.rangeselect.selection[0]
        larger = self.freq<self.rangeselect.selection[1]
        ndx = smaller*larger
        return ndx 
    
    def phase_center(self):
        if self.rangeselect.selection is not None:
            c0 = 299792458.0
    #        mask = mag>0.04
            mask = self._get_index()
            p1 = polyfit(self.freq[mask]*1e6, self.pha[mask], 1)
            y3 = polyval(p1, self.freq*1e6)
            self.status_txt = "%s\nPhase center = %5.2f m" % (self.status_txt, -p1[0]*c0/(2*pi))
        else:
            pass
#        mf = dphi/df
#        
#        cmed = 2300000.0
#        mk0 = c0*mf/(2*pi)
#        mk = cmed*mf/(2*pi)
        
#        cmed = c0/(sqrt(1.8))
#        m = (gradient(pha)/2.0)#*mag
#        from pylab import plot, show
#        plot(m)
#        show()
#        k = 2*pi/wl
#        phi = sum(y2)
#        pc = phi/k
#        print pc, wl, k, phi
        
if __name__ == "__main__":
    file = '20090212_stk01_1000ns_phase.h5'
    dir = r'Y:\Data2008_09\Hubra_proc\IPY'
#    file = 'data0.h5'
#    dir = r'Y:\Data2008_09\Wisdom\2312'
    tp = TraceViewer(file, dir)
    tp.configure_traits()
    